// Copyright (C) 2019 Alibaba Group Holding Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ===========================================================================

#include "postgres.h"

#include <omp.h>
#include "access/reloptions.h"
#include "catalog/index.h"
#include "miscadmin.h"
#include "storage/bufmgr.h"
#include "storage/freespace.h"
#include "storage/indexfsm.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
#include "utils/pase_hash_table.h"

#include "pase.h"

PG_MODULE_MAGIC;

relopt_kind hnsw_relopt_kind;
relopt_kind ivfflat_relopt_kind;

void
_PG_init(void) {
    int totalCoreNum;
    totalCoreNum = omp_get_num_procs();

    // hnsw options
    hnsw_relopt_kind = add_reloption_kind();
    add_int_reloption(hnsw_relopt_kind, "dim",
                      "vector dimension",
                      256, 8, 512);
    add_int_reloption(hnsw_relopt_kind, "base_nb_num",
                      "hnsw base_nb_num",
                      16, 5, 64);
    add_int_reloption(hnsw_relopt_kind, "ef_build",
                      "build queue",
                      40, 10, 320);
    add_int_reloption(hnsw_relopt_kind, "ef_search",
                      "vector dimension",
                      50, 10, 400);
    add_int_reloption(hnsw_relopt_kind, "base64_encoded",
                      "whether data base64 encoded",
                      0, 0, 1);
    // ivfflat options
    ivfflat_relopt_kind = add_reloption_kind();
    add_int_reloption(ivfflat_relopt_kind, "clustering_type",
                      "clustering type: 0 centroid_file, 1 inner clustering",
                      0, 0, 1);
    add_int_reloption(ivfflat_relopt_kind, "distance_type",
                      "distance metric type:0 l2, 1 inner proudct, 2 cosine",
                      0, 0, 2);
    add_int_reloption(ivfflat_relopt_kind, "dimension",
                      "vector dimension",
                      1, 1, 1024);
    add_int_reloption(ivfflat_relopt_kind, "open_omp",
                      "whether open omp",
                      0, 0, 1);
    add_int_reloption(ivfflat_relopt_kind, "omp_thread_num",
                      "omp thread number",
                      0, 1, totalCoreNum);
    add_int_reloption(ivfflat_relopt_kind, "base64_encoded",
                      "data whether base64 encoded",
                      0, 0, 1);
    add_string_reloption(ivfflat_relopt_kind, "clustering_params",
                      "clustering parameters", "", NULL);

}

/*
 * hnsw_options函数用于解析HNSW索引的选项参数
 *
 * 参数:
 *   reloptions: 索引的选项参数，类型为Datum
 *   validate: 是否验证选项参数，类型为bool
 *
 * 返回值:
 *   返回解析后的HNSWOptions结构体的字节表示(bytea *)
 */
bytea *
hnsw_options(Datum reloptions, bool validate) {
    HNSWOptions *opts;       // HNSWOptions结构体指针，用于存储解析后的选项参数
    relopt_value *options;   // 解析后的选项参数数组
    
    int numOptions;          // 解析后的选项参数数量

    opts = MakeDefaultHNSWOptions();  // 创建默认的HNSWOptions选项结构体

    // 定义选项参数的解析规则
    static const relopt_parse_elt hnsw_relopt_tab[] = {
        {"dim", RELOPT_TYPE_INT, offsetof(HNSWOptions, dim)},
        {"base_nb_num", RELOPT_TYPE_INT, offsetof(HNSWOptions, base_nb_num)},
        {"ef_build", RELOPT_TYPE_INT, offsetof(HNSWOptions, ef_build)},
        {"ef_search", RELOPT_TYPE_INT, offsetof(HNSWOptions, ef_search)},
        {"base64_encoded", RELOPT_TYPE_INT, offsetof(HNSWOptions, base64_encoded)}
    };

    // 解析选项参数并获取解析结果
    options = parseRelOptions(reloptions, validate, hnsw_relopt_kind, &numOptions);

    // 如果解析后的选项参数数量小于5，则报错
    if (numOptions < 5) {
        elog(ERROR, "options format error");
    }

    // 根据解析后的选项参数分配HNSWOptions结构体内存，并将解析结果填充到结构体中
    opts = allocateReloptStruct(sizeof(HNSWOptions), options, numOptions);
    fillRelOptions((void *)opts, sizeof(HNSWOptions), options, numOptions,
                   validate, hnsw_relopt_tab, lengthof(hnsw_relopt_tab));

    // 释放解析后的选项参数数组的内存
    pfree(options);

    // 返回HNSWOptions结构体的字节表示
    return (bytea *)opts;
}


PG_FUNCTION_INFO_V1(pase_ivfflat);

Datum
pase_ivfflat(PG_FUNCTION_ARGS) {
  IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);

  amroutine->amstrategies = 1;
  amroutine->amsupport = 1;
  amroutine->amcanorder = false;
  amroutine->amcanorderbyop = true;
  amroutine->amcanbackward = false;
  amroutine->amcanunique = false;
  amroutine->amcanmulticol = true;
  amroutine->amoptionalkey = true;
  amroutine->amsearcharray = false;
  amroutine->amsearchnulls = false;
  amroutine->amstorage = false;
  amroutine->amclusterable = false;
  amroutine->ampredlocks = false;
  amroutine->amcanparallel = false;
  amroutine->amkeytype = InvalidOid;

  amroutine->ambuild = ivfflat_build;
  amroutine->ambuildempty = ivfflat_buildempty;
  amroutine->aminsert = ivfflat_insert;
  amroutine->ambulkdelete = ivfflat_bulkdelete;
  amroutine->amvacuumcleanup = ivfflat_vacuumcleanup;
  amroutine->amcanreturn = NULL;
  amroutine->amcostestimate = ivfflat_costestimate;
  amroutine->amoptions = ivfflat_options;
  amroutine->amproperty = NULL;
  amroutine->ambeginscan = ivfflat_beginscan;
  amroutine->amrescan = ivfflat_rescan;
  amroutine->amgettuple = ivfflat_gettuple;
  amroutine->amgetbitmap = ivfflat_getbitmap;
  amroutine->amendscan = ivfflat_endscan;
  amroutine->ammarkpos = NULL;
  amroutine->amrestrpos = NULL;
  amroutine->amestimateparallelscan = NULL;
  amroutine->aminitparallelscan = NULL;
  amroutine->amparallelrescan = NULL;

  PG_RETURN_POINTER(amroutine);
}

/*
ambuild: 用于构建新的索引。
ambuildempty: 用于构建一个空的索引。
aminsert: 用于向索引中插入新的项。
ambulkdelete: 用于从索引中批量删除项。
amvacuumcleanup: 用于清理索引，通常在 VACUUM 操作后调用。
amcanreturn: 检查索引是否可以直接返回指定的列。
amcostestimate: 估计使用索引执行特定扫描的成本。
amoptions: 处理索引选项。
amproperty: 返回索引的特性或统计信息。
ambeginscan: 开始一个新的索引扫描。
amrescan: 重新扫描索引。
amgettuple: 从索引中获取下一个元组。
amgetbitmap: 从索引中获取一组元组。
amendscan: 结束索引扫描。
ammarkpos: 标记当前扫描位置。
amrestrpos: 重置扫描位置。

*/
PG_FUNCTION_INFO_V1(pase_hnsw);

Datum
pase_hnsw(PG_FUNCTION_ARGS) {
  IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);

  amroutine->amstrategies = 1; // 设置索引策略的数量
  amroutine->amsupport = 1; // 设置索引支持的操作符的数量
  amroutine->amcanorder = false; // 设置索引是否支持排序
  amroutine->amcanorderbyop = true; // 设置索引是否支持通过操作符进行排序
  amroutine->amcanbackward = false; // 设置索引是否支持反向扫描
  amroutine->amcanunique = false; // 设置索引是否支持唯一约束
  amroutine->amcanmulticol = true; // 设置索引是否支持多列索引
  amroutine->amoptionalkey = true; // 设置索引是否支持可选键
  amroutine->amsearcharray = false; // 设置索引是否支持数组搜索
  amroutine->amsearchnulls = false; // 设置索引是否支持空值搜索
  amroutine->amstorage = false; // 设置索引是否支持存储
  amroutine->amclusterable = false; // 设置索引是否支持聚簇
  amroutine->ampredlocks = false; // 设置索引是否支持预测锁
  amroutine->amcanparallel = false; // 设置索引是否支持并行查询
  amroutine->amkeytype = InvalidOid; // 设置索引键的类型


  amroutine->ambuild = hnsw_build;                 // ambuild: 用于构建新的索引。
  amroutine->ambuildempty = hnsw_buildempty;       // ambuildempty: 用于构建一个空的索引。
  amroutine->aminsert = hnsw_insert;               // aminsert: 用于向索引中插入新的项。
  amroutine->ambulkdelete = NULL;
  amroutine->amvacuumcleanup = hnsw_vacuumcleanup; // amvacuumcleanup: 用于清理索引，通常在 VACUUM 操作后调用
  amroutine->amcanreturn = NULL;
  amroutine->amcostestimate = hnsw_costestimate;   // amcostestimate: 估计使用索引执行特定扫描的成本。
  amroutine->amoptions = hnsw_options;             // amoptions: 处理索引选项。
  amroutine->amproperty = NULL;
  amroutine->ambeginscan = hnsw_beginscan;         // ambeginscan: 开始一个新的索引扫描。
  amroutine->amrescan = hnsw_rescan;               // amrescan: 重新扫描索引。
  amroutine->amgettuple = hnsw_gettuple;           // amgettuple: 从索引中获取下一个元组。
  amroutine->amgetbitmap = NULL;
  amroutine->amendscan = hnsw_endscan;             // amendscan: 结束索引扫描。
  amroutine->ammarkpos = NULL;
  amroutine->amrestrpos = NULL;
  amroutine->amestimateparallelscan = NULL;
  amroutine->aminitparallelscan = NULL;
  amroutine->amparallelrescan = NULL;

  PG_RETURN_POINTER(amroutine);
}
